import torch
from torch import nn
from torch.optim import Adam
from torch.utils.data import DataLoader
# 训练数据加载器
from mydatasets import my_dataset
from torch.utils.tensorboard import SummaryWriter
from vkmodel import vkmodel

if __name__ == '__main__':
    #test_dataset=my_dataset("./datasets/test/")
    # batch_size理解为超参数，多少个数据集为一组
    #test_dataloader=DataLoader(test_dataset,batch_size=40,shuffle=True)
    # print(torch.cuda.is_available())
    # print(torch.cuda.device_count())
    # 调用去学习
    train_dataset=my_dataset("./datasets/train/")
    train_dataloader=DataLoader(train_dataset,batch_size=40,shuffle=True)

#     vkmodel=vkmodel().cuda()
#     loss_fn=nn.MultiLabelSoftMarginLoss().cuda()
#     optim=Adam(vkmodel.parameters(),lr=0.001)
#     writer = SummaryWriter("logs")
#     total_step=0
#     for epoch in range(10):
#         print("外层训练次数{}".format(epoch))
#         for i, (images, lables) in enumerate(train_dataloader):
#             images=images.cuda()
#             lables=lables.cuda()
#             vkmodel.train()
#             ouput = vkmodel(images)
#             loss = loss_fn(ouput, lables)
#             optim.zero_grad()
#             loss.backward()
#             optim.step()
#             total_step+=1
#             if i % 100 == 0:
#                 print("训练次数{},损失率:{}".format(i, loss.item()))
#                 writer.add_scalar("loss",loss,total_step)
#
#
# torch.save(vkmodel,"model.pth")
        # torch.Size([40, 1, 60, 160])
        # torch.Size([图片数, 图片颜色类型, 宽度, 高度])